import functools
import gm_manager


def local_doc_content_crypto_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # 处理位置参数检查并加密内容
        new_args = []
        for arg in args:
            if isinstance(arg, dict) and 'content' in arg:
                en_content, iv = gm_manager.encrypt_local_document(arg['content'])
                arg['content'] = en_content
                arg['sm4_iv'] = iv
            new_args.append(arg)

        # 处理关键字参数检查并加密内容
        new_kwargs = {}
        for key, value in kwargs.items():
            if isinstance(value, dict) and 'content' in value:
                en_content, iv = gm_manager.encrypt_local_document(value['content'])
                value['content'] = en_content
                value['sm4_iv'] = iv
            new_kwargs[key] = value

        result = func(*new_args, **new_kwargs)

        # 处理输出检查（与原代码逻辑一致，保持不变）
        try:
            if isinstance(result, list):
                for element in result:
                    if isinstance(element, dict) and 'content' in element:
                        de_content = gm_manager.decrypt_local_document(element['content'], element['sm4_iv'])
                        element['content'] = de_content
            return result
        except Exception as e:
            print(f"出现错误: {e}")  # 打印错误信息，方便调试，也可以选择更合适的记录错误的方式
            return {"error": str(e)}  # 返回包含错误信息的字典，便于调用者统一处理错误情况
    return wrapper

def public_doc_content_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        # 处理位置参数检查
        for arg in args:
            if isinstance(arg, dict) and 'encrypted_data' in arg:
                # 加密和签名
                content_dict = arg['encrypted_data']
                encrypted_data, encrypted_sm4_key, sm4_iv, sm3_signature = gm_manager.encrypt_and_sign_document(content_dict,content_dict["receiving_unit"])
                arg['encrypted_data'] = encrypted_data
                arg['encrypted_sm4_key'] = encrypted_sm4_key
                arg['sm4_iv'] = sm4_iv
                arg['sm3_signature'] = sm3_signature

        # 处理关键字参数检查
        for value in kwargs.values():
            if isinstance(value, dict) and 'encrypted_data' in value:
                content_dict = value['encrypted_data']
                encrypted_data, encrypted_sm4_key, sm4_iv, sm3_signature = gm_manager.encrypt_and_sign_document(content_dict,content_dict["receiving_unit"])
                value['encrypted_data'] = encrypted_data
                value['encrypted_sm4_key'] = encrypted_sm4_key
                value['sm4_iv'] = sm4_iv
                value['sm3_signature'] = sm3_signature

        result = func(*args, **kwargs)

        # 处理输出检查
        if isinstance(result, (list, tuple)):
            for element in result:
                if isinstance(element, dict) and 'encrypted_data' in element:
                    # 解密与验证
                    element['is_verified'] = False
                    encrypted_dict = element['encrypted_data']
                    decrypted_dict, verification_result = gm_manager.decrypt_and_verify_document(encrypted_dict, element['receiving_unit'],element['encrypted_sm4_key'], element['sm4_iv'], element['sm3_signature'])
                    element['encrypted_data'] = decrypted_dict
                    element['is_verified'] = verification_result

        return result
    return wrapper
